NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR NNX migration prep (2/N): NNX utils and sharding utilities #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop #3500)
4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR [NNX] NNX migration prep (4/N): sharding tools and post-training fixes #3652)
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
6. ✅ NNX-native DPO.
7. ✅ NNX-native MaxEngine inference. (PR [NNX] NNX migration prep (7/N): NNX-native MaxEngine inference #3821)
8. ✅ NNX-native LoRA + GRPO. (PR [NNX] NNX migration prep (8/N): NNX native lora grpo #3824)
9. ✅ NNX-aware QK-Clip + remaining checkpoint utilities. (PR [NNX] NNX migration prep (9/N): NNX-aware QK-Clip + checkpoint utilities #3836)
   9.5. ✅ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix. (PR [NNX] NNX migration prep (9.5/N): NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix #3844)
10. ✅ Vocab tiling custom_vjp for NNX.
11. ✅ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations. (PR [NNX] NNX migration (11/N): set pure_nnx / enable_nnx / pure_nnx_decoder defaults to True #3526)
12. Delete Linen-specific code paths and NNX compatibility flags — delivered as 5 stacked PRs (12.1–12.5):
    12.1. ✅ Collapse Linen dispatch in utils / quantization / checkpointing / sharding. (PR [NNX] Delete Linen 1/5: collapse dispatch in utils, quantization, checkpointing, sharding #4357)
    12.2. ✅ Collapse Linen dispatch in trainers (pre-train, DiLoCo, GRPO); regenerate the NNX DPO golden. (PR [NNX] Delete Linen 2/5: collapse dispatch in trainers (pre-train, DiLoCo, GRPO) #4358)
    12.3. ✅ Collapse Linen dispatch in inference (maxengine, kvcache, vLLM, LoRA). (PR [NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA) #4359)
    12.4. ✅ Delete the Linen decoder/attention/embedding layers + *_as_linen model wrappers. (PR [NNX] Delete Linen 4/5: remove the Linen decoder/attention layers and *_as_linen model wrappers #4360)
    12.5. 🔄 [This PR] Remove the pure_nnx / enable_nnx / pure_nnx_decoder flags from types.py (Fields + dead qwix validator), base.yml, inference/vllm.yml, the distillation configs, and pyconfig_deprecated.py, plus doc/notebook/script mentions. No source or test reads the flags anymore; config init succeeds with them absent. Completes the Linen deletion (≈ −5.3k lines across the series); the original single-branch PR [NNX] NNX migration (12/N): delete Linen code paths, classes, and NNX compatibility flags #4038 can be closed. Stacks on 12.4.

Description

Final PR of the stacked "Delete Linen" series (see 1/5). With every read gone after 1–4, this removes the three migration flags — pure_nnx, enable_nnx, pure_nnx_decoder — from types.py (the Fields plus the now-dead qwix validator), base.yml, inference/vllm.yml, the distillation configs, and pyconfig_deprecated.py, and clears the remaining mentions in the docs, the LoRA demo notebook, and the distillation / gemma3-LoRA shell scripts.

Nothing reads the flags anymore (grep -rE 'pure_nnx|enable_nnx|pure_nnx_decoder' src/ tests/, excluding reference_hlo, is empty), and config init succeeds with them absent, including under use_qwix_quantization=True. After this the Linen deletion is complete — roughly −5.3k lines across the series — and the original single-branch PR #4038 can be closed. Base branch is feat/nnx-del-linen-4-layers. ~14 files, +9 / −54.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.